import os

import pygame
import torch
import network
from MCTS import MCTS


def main():
    oppo = "model_80.pkl"
    board_size = 15
    try:
        Net = torch.load(oppo, map_location="cpu")
        Net.device = "cpu"
        tree = MCTS(board_size=board_size, neural_network=Net)
    except Exception as e:
        print(e)
        raise ValueError("The parameter oppo must be a pretrained model")

    GRID_WIDTH = 36
    WIDTH = (board_size + 2) * GRID_WIDTH
    HEIGHT = (board_size + 2) * GRID_WIDTH
    FPS = 30
    record, game_continue = tree.interact_game_init()

    # 下面这些都是跟游戏相关的
    WHITE = (255, 255, 255)
    BLACK = (0, 0, 0)
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("五子棋")
    clock = pygame.time.Clock()
    base_folder = os.path.dirname(__file__)
    img_folder = os.path.join(base_folder, "images")
    background_img = pygame.image.load(os.path.join(img_folder, "back.png")).convert()
    background = pygame.transform.scale(background_img, (WIDTH, HEIGHT))
    back_rect = background.get_rect()

    def draw_background(surf):
        screen.blit(background, back_rect)
        rect_lines = [
            ((GRID_WIDTH, GRID_WIDTH), (GRID_WIDTH, HEIGHT - GRID_WIDTH)),
            ((GRID_WIDTH, GRID_WIDTH), (WIDTH - GRID_WIDTH, GRID_WIDTH)),
            ((GRID_WIDTH, HEIGHT - GRID_WIDTH), (WIDTH - GRID_WIDTH, HEIGHT - GRID_WIDTH)),
            ((WIDTH - GRID_WIDTH, GRID_WIDTH), (WIDTH - GRID_WIDTH, HEIGHT - GRID_WIDTH)),
        ]
        for line in rect_lines:
            pygame.draw.line(surf, BLACK, line[0], line[1], 2)

        for i in range(board_size):
            pygame.draw.line(surf, BLACK, (GRID_WIDTH * (2 + i), GRID_WIDTH), (GRID_WIDTH * (2 + i), HEIGHT - GRID_WIDTH))
            pygame.draw.line(surf, BLACK, (GRID_WIDTH, GRID_WIDTH * (2 + i)), (HEIGHT - GRID_WIDTH, GRID_WIDTH * (2 + i)))

        circle_center = [
            (GRID_WIDTH * 4, GRID_WIDTH * 4),
            (WIDTH - GRID_WIDTH * 4, GRID_WIDTH * 4),
            (WIDTH - GRID_WIDTH * 4, HEIGHT - GRID_WIDTH * 4),
            (GRID_WIDTH * 4, HEIGHT - GRID_WIDTH * 4),
        ]
        for cc in circle_center:
            pygame.draw.circle(surf, BLACK, cc, 5)

    running = True

    def draw_stone(scene):
        for i in range(board_size):
            for j in range(board_size):
                if record[i, j] == 1:
                    pygame.draw.circle(scene, BLACK, (int((i + 1.5) * GRID_WIDTH), int((j + 1.5) * GRID_WIDTH)), 16)
                elif record[i, j] == -1:
                    pygame.draw.circle(scene, WHITE, (int((i + 1.5) * GRID_WIDTH), int((j + 1.5) * GRID_WIDTH)), 16)

    while running:
        clock.tick(FPS)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if game_continue:
                    pos = event.pos
                    if pos[0] < GRID_WIDTH or pos[1] < GRID_WIDTH or pos[0] > WIDTH - GRID_WIDTH or pos[1] > HEIGHT - GRID_WIDTH:
                        pass
                    else:
                        grid = (int((pos[0] - GRID_WIDTH) / GRID_WIDTH), int((pos[1] - GRID_WIDTH) / GRID_WIDTH))
                        record, game_continue = tree.interact_game1(grid)
                        draw_background(screen)
                        draw_stone(screen)
                        pygame.display.flip()
                        record, game_continue = tree.interact_game2(grid, game_continue, record)
                else:
                    running = False
        draw_background(screen)
        draw_stone(screen)
        pygame.display.flip()
    pygame.quit()


if __name__ == "__main__":
    main()
